use anyhow::Result;
use serde_json::{json, Value};
use std::collections::HashMap;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tracing::{info, error};

pub struct SimpleMcpServer {
    name: String,
    version: String,
    tools: HashMap<String, ToolDefinition>,
    tool_handler: Option<Box<dyn Fn(ToolCall) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<ToolResult>> + Send>> + Send + Sync>>,
}

#[derive(Debug, Clone)]
pub struct ToolDefinition {
    pub name: String,
    pub description: String,
    pub schema: Value,
}

#[derive(Debug, Clone)]
pub struct ToolCall {
    pub name: String,
    pub arguments: Option<Value>,
}

#[derive(Debug, Clone)]
pub struct ToolResult {
    pub content: Vec<Content>,
    pub is_error: bool,
}

#[derive(Debug, Clone)]
pub struct Content {
    pub content_type: String,
    pub text: String,
}

impl ToolResult {
    pub fn success(data: Value) -> Self {
        Self {
            content: vec![Content {
                content_type: "text".to_string(),
                text: data.to_string(),
            }],
            is_error: false,
        }
    }
    
    pub fn error(message: &str) -> Self {
        Self {
            content: vec![Content {
                content_type: "text".to_string(),
                text: message.to_string(),
            }],
            is_error: true,
        }
    }
    
    pub fn is_error(&self) -> bool {
        self.is_error
    }
}

impl SimpleMcpServer {
    pub fn new(name: &str, version: &str) -> Self {
        Self {
            name: name.to_string(),
            version: version.to_string(),
            tools: HashMap::new(),
            tool_handler: None,
        }
    }
    
    pub fn with_tool(mut self, name: &str, description: &str, schema: Value) -> Self {
        self.tools.insert(name.to_string(), ToolDefinition {
            name: name.to_string(),
            description: description.to_string(),
            schema,
        });
        self
    }
    
    pub fn set_tool_handler<F, Fut>(&mut self, handler: F)
    where
        F: Fn(ToolCall) -> Fut + Send + Sync + 'static,
        Fut: std::future::Future<Output = Result<ToolResult>> + Send + 'static,
    {
        self.tool_handler = Some(Box::new(move |call| {
            Box::pin(handler(call))
        }));
    }
    
    pub async fn run(&self) -> Result<()> {
        info!("Starting MCP Server: {} v{}", self.name, self.version);
        info!("Available tools: {:?}", self.tools.keys().collect::<Vec<_>>());
        
        let stdin = tokio::io::stdin();
        let mut reader = BufReader::new(stdin);
        let mut stdout = tokio::io::stdout();
        
        // Send server info
        let server_info = json!({
            "jsonrpc": "2.0",
            "result": {
                "protocolVersion": "2025-06-18",
                "capabilities": {
                    "tools": {}
                },
                "serverInfo": {
                    "name": self.name,
                    "version": self.version
                }
            }
        });
        
        stdout.write_all(format!("{}\n", server_info).as_bytes()).await?;
        stdout.flush().await?;
        
        let mut line = String::new();
        loop {
            line.clear();
            match reader.read_line(&mut line).await {
                Ok(0) => break, // EOF
                Ok(_) => {
                    if let Ok(request) = serde_json::from_str::<Value>(&line) {
                        if let Some(response) = self.handle_request(request).await {
                            stdout.write_all(format!("{}\n", response).as_bytes()).await?;
                            stdout.flush().await?;
                        }
                    }
                }
                Err(e) => {
                    error!("Error reading from stdin: {}", e);
                    break;
                }
            }
        }
        
        Ok(())
    }
    
    async fn handle_request(&self, request: Value) -> Option<Value> {
        let method = request.get("method")?.as_str()?;
        let id = request.get("id");
        
        match method {
            "tools/list" => {
                let tools: Vec<_> = self.tools.values().map(|tool| {
                    json!({
                        "name": tool.name,
                        "description": tool.description,
                        "inputSchema": tool.schema
                    })
                }).collect();
                
                Some(json!({
                    "jsonrpc": "2.0",
                    "id": id,
                    "result": {
                        "tools": tools
                    }
                }))
            }
            "tools/call" => {
                if let Some(handler) = &self.tool_handler {
                    let params = request.get("params")?;
                    let name = params.get("name")?.as_str()?.to_string();
                    let arguments = params.get("arguments").cloned();
                    
                    let call = ToolCall { name, arguments };
                    
                    match handler(call).await {
                        Ok(result) => {
                            Some(json!({
                                "jsonrpc": "2.0",
                                "id": id,
                                "result": {
                                    "content": result.content.iter().map(|c| {
                                        json!({
                                            "type": c.content_type,
                                            "text": c.text
                                        })
                                    }).collect::<Vec<_>>(),
                                    "isError": result.is_error
                                }
                            }))
                        }
                        Err(e) => {
                            Some(json!({
                                "jsonrpc": "2.0",
                                "id": id,
                                "error": {
                                    "code": -32603,
                                    "message": format!("Tool execution failed: {}", e)
                                }
                            }))
                        }
                    }
                } else {
                    Some(json!({
                        "jsonrpc": "2.0",
                        "id": id,
                        "error": {
                            "code": -32601,
                            "message": "No tool handler configured"
                        }
                    }))
                }
            }
            _ => {
                Some(json!({
                    "jsonrpc": "2.0",
                    "id": id,
                    "error": {
                        "code": -32601,
                        "message": format!("Method not found: {}", method)
                    }
                }))
            }
        }
    }
}